#### Simulate BR term structures for each step from 1% to 5% in increments of 0.5%

# Simulate rates out tt=500 years
N=10e3
tt=500

#### Sigma's posterior quantiles, from Michael Bauer:
#     sigma_u:
#  1st Qu.  Median   3rd Qu.  
#   0.1946  0.2039   0.2140  
set.seed(1)
# Use inverse gamma distribution (the BR prior distribution), calibrated to match
# the provided posterior quantiles
adj.factor = (1.003*0.2047/0.10)^2 # this adjusts the inverse gamma to match the BR posterior
sig = sqrt(1/rgamma(N, shape=100/2, scale=0.04*(100+2)/2)*adj.factor)

quantile(sig, probs = c(0.25, 0.5, 0.75)) # should come close to 0.1946, 0.2039, 0.2140

# rescale so it's not in percentage points anymore
sig = sig/100 

# Generate innovations/shocks
set.seed(1)
head(rnorm(tt*N, 0, sig)) # check values

eps = matrix(NA, nrow=tt, ncol=N)
for (t in 1:tt) eps[t,] = rnorm(N, 0, sig)
eps[1,] = 0 # first period has no shock to fix starting point

# Accumuated shocks
eps.cum = apply(eps, 2, cumsum)

# Set up array to store results
ce.store = array(NA, dim=c(tt, 8),
                 dimnames=list(1:tt,
                               c('1.5%','2%','2.5%','3%','3.5%','4%','4.5%','5%')))
r0s = seq(0.015, 0.05, by=0.005)

for (i in 1:length(r0s)) { # for each starting rate...
  # ...simulate real rate, constrained st avg rate>0 ----
  r0 = r0s[i] # starting rate
  
  r = r0 + eps.cum
  
  # We're imposing the constraint that average rates cannot be negative.
  # This is the same as saying the discount factor cannot exceed 1, DF>=1, 
  # so that we get r^ce_t = log(DF)/-t >= 0
  # Do this by calculating the DF, then flooring it:
  DF = apply(r, 2, function(x) exp(-cumsum(x)))
  
  DF = pmin(DF, 1)
  
  # Check:
  # Check forward rates:
  r.floor = matrix(NA, nrow=nrow(DF), ncol=ncol(DF))
  r.floor[1,] = r[1,]
  r.floor[2:tt,] = - (log(DF[2:tt,]) - log(DF[1:(tt-1),]))
  
  # Check time average rates: raw r versus floored
  r.avg = apply(r, 2, function(x) cumsum(x)/seq_along(x))
  r.floor.avg = apply(r.floor, 2, function(x) cumsum(x)/seq_along(x))
  
  range(r); range(r.floor) # note: forward rates can be negative, even after flooring...
  range(r.avg); range(r.floor.avg) # ...but time average rates are >=0 after flooring
  
  E.DF = apply(DF, 1, mean)
  
  ce = log(E.DF)/-c(1:tt)
  
  ce.store[,i] = ce
}

ce.BR = ce.store
save('ce.BR', file=data.dir%&%'/extra_BR_term_structures.RData')

ce.BR = ce.BR[-1,] # drop year 1 which is the near-term rate by construction
# Calibrate
output.BR <- array(list(), dim=c(dim(ce.BR)[2]), 
                   dimnames=dimnames(ce.BR)[2])

set.seed(1)

library(progress)
pb = progress_bar$new(total=length(output.BR))

n.yrs = 10
st = Sys.time()
for (r in 1:length(output.BR)) {
  if (any(is.na(ce.BR[,r]))) next # don't have all rates for Groom or Freeman
  # Solve for optimal eta
  sol.temp = nlm(NP_MWS_ce_dist, target.ce=ce.BR[,r], 1.3, 
                 n.yrs=n.yrs, iterlim=1000, gradtol=1e-6,
                 g=global_growth_percap)
  
  if (sol.temp$code>=4) {message(paste(r,' failed')); next}
  
  # Output term structure for optimal eta
  ce.fit = calc_scc_ce(sol.temp$estimate, target.ce=ce.BR[,r], n.yrs=n.yrs,
                       g=global_growth_percap)
  
  # Find eta when rho is truncated at zero
  if (ce.fit$rho<0) {
    unconstrained.rho = ce.fit$rho
    unconstrained.eta = ce.fit$eta
    
    eta.grid.temp = seq(0,2, by=0.01)
    # since near-term g is ~1.5%, eta=r/g will be very close to 1.5/1.6~=1 for r=1.5%, and close to 0.66% for r=1%, since near-term g is
    
    near.term.diff = function(eta) {
      near.term.g.ce = log(apply(exp(-eta*apply(global_growth_percap[1:n.yrs,], 2, cumsum)), 1, mean)) * -1/c(1:n.yrs)
      return(mean((near.term.g.ce - ce.BR[1:n.yrs,r])^2))
    }
    ntd = sapply(eta.grid.temp, near.term.diff)
    cbind(eta.grid.temp, ntd)
    plot(ntd ~ eta.grid.temp); abline(h=0)
    eta.opt = eta.grid.temp[which.min(ntd)]
    
    ce.fit = calc_scc_ce(eta.opt,rho=0, target.ce=ce.BR[,r], n.yrs=n.yrs,
                         g=global_growth_percap)
    ce.fit$unconstrained.rho = unconstrained.rho
    ce.fit$unconstrained.eta = unconstrained.eta
    
  }
  
  ce.fit$sse = sol.temp$`minimum` # save sum of squared errors
  
  output.BR[[r]] = ce.fit
  pb$tick()
  
}
rm(list=c('DF','eps','eps.cum','r.avg','r.floor','r.floor.avg'))

ed = Sys.time()
difftime(ed, st)

## Plot BR term structures
pdf(figure.dir%&%'/BauerRudebusch_RhoEta_Estimates_AllRates_'%&%today()%&%'.pdf',
    family='serif', width=7*1.5, height=7)
dimnames(output)

plot.ces(output.BR['5%'][[1]], col=3, main='Bauer Rudebusch (2020) Fits', maxbuff=0.005)# 3=red
plot.ces(output.BR['4.5%'][[1]], addline=TRUE, col=7, cex.axis=1.5, cex.lab=2) # 7=orange
plot.ces(output.BR['4%'][[1]], addline=TRUE, col=8, cex.axis=1.5, cex.lab=2) # 8=yellow
plot.ces(output.BR['3.5%'][[1]], addline=TRUE, col=5, cex.axis=1.5, cex.lab=2) # 5=brown
plot.ces(output.BR['3%'][[1]], addline=TRUE, col=1, cex.axis=1.5, cex.lab=2) # 1=black
plot.ces(output.BR['2.5%'][[1]], addline=TRUE, col=colorspace::darken(9, 0.4), cex.axis=1.5, cex.lab=2) # 9=light gray
plot.ces(output.BR['2%'][[1]], addline=TRUE, col=4, cex.axis=1.5, cex.lab=2) # 4=green
plot.ces(output.BR['1.5%'][[1]], addline=TRUE, col=2, cex.axis=1.5, cex.lab=2) # 2=blue
abline(h=0)

legend.values = sapply(output.BR, function(x) c(rho=x$rho, eta=x$eta))
constr.rho = which(legend.values['rho',]==0)
legend.values['rho',] = round(legend.values['rho',],3)*100
legend.values['eta',] = round(legend.values['eta',],2)
legend.values['rho',] = paste0(legend.values['rho',],'%')
legend.values[,constr.rho] = paste0(legend.values[,constr.rho],'*')

legend('top', lty=2:1, lwd=3, bg='white', legend=c('Fitted Rate', 'Target Rate'), cex=1.3, col='black')
legend('topright', fill=c(3,7,8,5,1,colorspace::darken(9, 0.4),4,2), border=NA, bg='white', cex=1.2,
       legend = c(
         as.expression(bquote(
           paste(r[0], ' = 5.0%  ',
                 rho,' = ', .(legend.values['rho','5%']), '    ',
                 eta,' = ', .(legend.values['eta','5%'])))),
         as.expression(bquote(
           paste(r[0], ' = 4.5%  ',
                 rho,' = ', .(legend.values['rho','4.5%']), '    ',
                 eta,' = ', .(legend.values['eta','4.5%'])))),
         as.expression(bquote(
           paste(r[0], ' = 4.0%  ',
                 rho,' = ', .(legend.values['rho','4%']), '    ',
                 eta,' = ', .(legend.values['eta','4%'])))),
         as.expression(bquote(
           paste(r[0], ' = 3.5%  ',
                 rho,' = ', .(legend.values['rho','3.5%']), '    ',
                 eta,' = ', .(legend.values['eta','3.5%'])))),
         as.expression(bquote(
           paste(r[0], ' = 3.0%  ',
                 rho,' = ', .(legend.values['rho','3%']), '    ',
                 eta,' = ', .(legend.values['eta','3%'])))),
         as.expression(bquote(
           paste(r[0], ' = 2.5%  ',
                 rho,' = ', .(legend.values['rho','2.5%']), '    ',
                 eta,' = ', .(legend.values['eta','2.5%'])))),
         as.expression(bquote(
           paste(r[0], ' = 2.0%  ',
                 rho,' = ', .(legend.values['rho','2%']), '    ',
                 eta,' = ', .(legend.values['eta','2%'])))),
         as.expression(bquote(
           paste(r[0], ' = 1.5%  ',
                 rho,' = ', .(legend.values['rho','1.5%']), '    ',
                 eta,' = ', .(legend.values['eta','1.5%']))))
       )
)
dev.off()

